import numpy as np
import matplotlib.pyplot as plt

def aceleracao_gravitacional(r, G, M, m):
    """
    Calcula a aceleração gravitacional entre dois corpos usando a Lei da Gravitação Universal.

    Parâmetros:
        r (np.ndarray): vetor de posição do corpo (2D).
        G (float): constante gravitacional no sistema de unidades adotado.
        M (float): massa do corpo central (ex: Sol).
        m (float): massa do corpo em movimento (ex: Terra).

    Retorna:
        np.ndarray: vetor aceleração (2D).
    """
    distancia = np.linalg.norm(r)
    Fg = -G * M * m * r / distancia**3
    return Fg / m


def metodo_euler_orbita(r0, v0, G, M, m, t0, tf, dt):
    """
    Simula a órbita de um corpo ao redor de outro usando o Método de Euler.

    Parâmetros:
        r0 (np.ndarray): vetor posição inicial (2D).
        v0 (np.ndarray): vetor velocidade inicial (2D).
        G (float): constante gravitacional.
        M (float): massa do corpo central.
        m (float): massa do corpo em órbita.
        t0 (float): tempo inicial.
        tf (float): tempo final.
        dt (float): passo de tempo.

    Retorna:
        tuple: (array de tempos, array de posições, array de velocidades)
    """
    n = int((tf - t0) / dt)
    t = np.linspace(t0, tf, n)
    r = np.empty((n, 2))
    v = np.empty((n, 2))

    r[0] = r0.copy()
    v[0] = v0.copy()

    for i in range(n - 1):
        a = aceleracao_gravitacional(r[i], G, M, m)
        r[i+1] = r[i] + v[i] * dt
        v[i+1] = v[i] + a * dt

    return t, r, v


def plotar_orbita(r):
    """
    Plota a órbita de um corpo com base nas posições calculadas.

    Parâmetros:
        r (np.ndarray): array de posições (n x 2).
    """
    plt.figure()
    plt.plot(r[:, 0], r[:, 1], label="Órbita (Euler)")
    plt.xlabel("x (UA)")
    plt.ylabel("y (UA)")
    plt.title("Simulação da Órbita da Terra em torno do Sol")
    plt.grid()
    plt.axis('equal')
    plt.legend()
    plt.show()


# ======================
# EXEMPLO DE USO
# ======================

if __name__ == "__main__":
    # Sistema astronômico de unidades:
    # G em AU³/(M*ano²), M em massas solares, tempo em anos

    G = 4 * np.pi**2  # AU³ / (M * ano²)
    M = 1.0           # Massa do Sol (em unidades solares)
    m = 1.0           # Massa da Terra (não afeta o movimento aqui)

    r0 = np.array([1.0, 0.0])  # 1 AU de distância do Sol
    v0 = np.array([0.0, 2 * np.pi])  # Velocidade para 1 ano orbital (UA/ano)

    t0 = 0.0
    tf = 4.0  # Simular 4 anos
    dt = 0.01

    t, r, v = metodo_euler_orbita(r0, v0, G, M, m, t0, tf, dt)
    plotar_orbita(r)